"""Train the model"""

import argparse
import logging
import os

import numpy as np
import torch
import torch.optim as optim
import yaml
from torch.autograd import Variable
from tqdm import tqdm

import utils
import model.mynet as net
import data.data_loader as data_loader
# from evaluate import evaluate

parser = argparse.ArgumentParser()
parser.add_argument('--model_dir', default='experiments/yaml_config',
                    help="Directory containing config.yaml")


def one_epoch(model, optimizer, loss_fn, dataloader, metrics, config, is_train):
    """Train the model on `num_steps` batches

    Args:
        model: (torch.nn.Module) the neural network
        optimizer: (torch.optim) optimizer for parameters of model
        loss_fn: a function that takes batch_output and batch_labels and computes the loss for the batch
        dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches training data
        metrics: (dict) a dictionary of functions that compute a metric using the output and labels of each batch
        config: (Params) hyperparameters from YAML config
        num_steps: (int) number of batches to train on, each of size params.batch_size
        is_train: (bool) whether or not the model is being trained
    """

    # set model training mode
    if is_train:
        model.train()
        desc = 'Train:'
    else:
        model.eval()
        desc = 'Valid:'

    # summary for current training loop and a running average object for loss
    summ = []
    loss_avg = utils.RunningAverage()

    # Use tqdm for progress bar
    with tqdm(total=len(dataloader), desc=desc) as t:
        for i, (train_batch, labels_batch) in enumerate(dataloader):
            # move to GPU if available
            if config['device']['gpu'] and torch.cuda.is_available():
                train_batch, labels_batch = train_batch.cuda(non_blocking=True), labels_batch.cuda(non_blocking=True)
            # convert to torch Variables
            train_batch, labels_batch = Variable(train_batch), Variable(labels_batch)

            # compute model output and loss
            output_batch = model(train_batch)
            loss = loss_fn(output_batch, labels_batch)

            # Evaluate summaries only once in a while
            if is_train:
                # clear previous gradients, compute gradients of all variables wrt loss
                optimizer.zero_grad()
                loss.backward()
                # performs updates using calculated gradients
                optimizer.step()

            if not is_train or i % config.get('save_summary_steps', 10) != 0:
                # extract data from torch Variable, move to cpu, convert to numpy arrays
                output_batch = output_batch.data.cpu().numpy()
                labels_batch = labels_batch.data.cpu().numpy()

                # compute all metrics on this batch
                summary_batch = {metric: metrics[metric](output_batch, labels_batch)
                                 for metric in metrics}
                summary_batch['loss'] = loss.item()
                summ.append(summary_batch)

            # update the average loss
            loss_avg.update(loss.item())

            t.set_postfix(loss='{:05.3f}'.format(loss_avg()))
            t.update()

    # compute mean of all metrics in summary
    metrics_mean = {metric: np.mean([x[metric]
                                     for x in summ]) for metric in summ[0]}
    metrics_string = " ; ".join("{}: {:05.3f}".format(k, v)
                                for k, v in metrics_mean.items())
    logging.info(f"- {desc} metrics: " + metrics_string)
    return metrics_mean


def train_and_evaluate(model, train_dataloader, val_dataloader,
                       optimizer, loss_fn, metrics,
                       config, model_dir, restore_file=None, patience=10):
    """Train the model and evaluate every epoch.

    Args:
        model: (torch.nn.Module) the neural network
        train_dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches training data
        val_dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches validation data
        optimizer: (torch.optim) optimizer for parameters of model
        loss_fn: a function that takes batch_output and batch_labels and computes the loss for the batch
        metrics: (dict) a dictionary of functions that compute a metric using the output and labels of each batch
        config: (Params) hyperparameters from YAML config
        model_dir: (string) directory containing config, weights and log
        restore_file: (string) optional- name of file to restore from (without its extension .pth.tar)
        patience: (int) number of epochs to wait for improvement before early stopping
    """
    # reload weights from restore_file if specified
    if restore_file is not None:
        restore_path = os.path.join(
            model_dir, restore_file + '.pth.tar')
        logging.info("Restoring parameters from {}".format(restore_path))
        utils.load_checkpoint(restore_path, model, optimizer)



    # Learning rate scheduler
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode='max',  # Reduce LR when accuracy stops decreasing
        factor=0.5,  # Reduce LR by half
        patience=3,  # Wait 5 epochs before reducing
        verbose=True
    )

    # Early stopping variables
    best_val_acc = 0.0  # Higher accuracy is better
    epochs_no_improve = 0
    best_epoch = 0

    for epoch in range(config['optim']['num_epochs']):
        # Run one epoch
        logging.info("Epoch {}/{}".format(epoch + 1, config['optim']['num_epochs']))

        # compute number of batches in one epoch (one full pass over the training set)
        one_epoch(model, optimizer, loss_fn, train_dataloader, metrics, config, True)

        # Evaluate for one epoch on validation set
        val_metrics = one_epoch(model,None, loss_fn, val_dataloader, metrics, config, False)

        val_acc = val_metrics['accuracy']
        is_best = val_acc >= best_val_acc

        # Learning rate scheduling
        scheduler.step(best_val_acc)
        current_lr = optimizer.param_groups[0]['lr']
        logging.info(f"Current learning rate: {current_lr:.6f}")

        # Save weights
        state = {
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'optim_dict': optimizer.state_dict(),
        }
        utils.save_checkpoint(state,
                              is_best=is_best,
                              checkpoint=model_dir)

        # If best_eval, best_save_path
        if is_best:
            logging.info(f"- Found new best accuracy: {val_acc:.3f}")
            best_val_acc = val_acc
            best_epoch = epoch + 1
            epochs_no_improve = 0
            # Save best val metrics in a json file in the model directory
            best_json_path = os.path.join(
                model_dir, "metrics_val_best_weights.json")
            utils.save_dict_to_json(val_metrics, best_json_path)
        else:
            epochs_no_improve += 1
            logging.info(f"- No improvement. Epochs without improvement: {epochs_no_improve}/{patience}")

        # Save latest val metrics in a json file in the model directory
        last_json_path = os.path.join(
            model_dir, "metrics_val_last_weights.json")
        utils.save_dict_to_json(val_metrics, last_json_path)

        # Early stopping check
        if epochs_no_improve >= patience:
            logging.info(
                f"Early stopping triggered after {epoch + 1} epochs. Best Acc: {best_val_acc:.3f} at epoch {best_epoch}")
            break
    logging.info(f"Training completed. Best validation Acc: {best_val_acc:.3f} at epoch {best_epoch}")

if __name__ == '__main__':

    # Load the parameters from json file
    args = parser.parse_args()

    # Set the logger
    os.makedirs(args.model_dir, exist_ok=True)
    utils.set_logger(os.path.join(args.model_dir, 'train.log'))

    # Check for existing best checkpoint
    best_checkpoint = os.path.join(args.model_dir, 'best.pth.tar')
    restore_file = 'best' if os.path.exists(best_checkpoint) else None

    yaml_path = os.path.join(args.model_dir, 'config.yaml')
    # print(yaml_path)
    with open(yaml_path, "r", encoding='utf-8') as f:
        config = yaml.safe_load(f)
    assert os.path.isfile(
        yaml_path), "No yaml configuration file found at {}".format(yaml_path)

    # use GPU if available
    device = torch.device(f"cuda:{config['device']['num']}" if torch.cuda.is_available() and config['device']['gpu'] else "cpu")
    logging.info(f'Using device: {device}')
    # Set the random seed for reproducible experiments
    torch.manual_seed(42)
    if config['device']['gpu']:
        torch.cuda.manual_seed(42)

    # Create the input data pipeline
    logging.info("Loading the datasets...")
    # fetch dataloaders
    dataloaders = data_loader.fetch_dataloader(
        ['train', 'val'], config)
    train_dl = dataloaders['train']
    val_dl = dataloaders['val']

    logging.info("- done.")

    # Define the model and optimizer
    model = net.Net(config).to(device)
    # model = net.Net(config).cuda() if config['device']['gpu'] else net.Net(config)
    optimizer = optim.Adam(model.parameters(),
                           lr=config['optim']['learning_rate'])

    # fetch loss function and metrics
    loss_fn = net.loss_fn
    metrics = net.metrics

    # Train the model
    logging.info("Starting training for {} epoch(s)".format(config['optim']['num_epochs']))
    train_and_evaluate(
        model=model,
        train_dataloader=train_dl,
        val_dataloader=val_dl,
        optimizer=optimizer,
        loss_fn=net.loss_fn,
        metrics=net.metrics,
        config=config,
        model_dir=args.model_dir,
        restore_file=restore_file
    )
